# 汇总报表的脚本

# 1. 生成某个月的表格模板
# 2. 发下去填写
# 3. 收集所有表格后，放到一个文件夹中
# 4. 将表格模板与每一个表格进行比较，凡是不相同的地方，都用红色标记出来，每个表格都与表格模板进行比较，就把数据收集完了
# 5. 生成标记了修改处的新表格
# 
# 关键参数：
# - 数据起始行
# - 最大列


import os
import shutil
import sys

from openpyxl import load_workbook
from openpyxl.styles import Font


START_ROW = 6
MAX_COL = 21
PRIMARY_COL = 21

def try_file_path(file_path):
    """测试是否存在文件，如果存在，返回带括号的文件名"""
    tmp_path = file_path
    (root, ext) = os.path.splitext(file_path)
    i = 1
    while os.path.exists(tmp_path):
        tmp_path = f"{root}({i}){ext}"
        i += 1

    return tmp_path


class CollectTable:
    def __init__(self, file_path, primary_col, start_row, max_col) -> None:
        self.file_path = file_path
        self.primary_col = primary_col
        self.start_row = start_row
        self.max_col = max_col
        self.wb = load_workbook(file_path)
        self.ws = self.wb.active
        pass

    def get_new_month(self, month):
        '''根据实际，修改表格模板'''

        # 修改月份标记
        self.ws['A2'].value = f'云南省各州（市）2024年能源外工业项目投资情况调度表（{month}月）'
        self.ws['L5'].value = f'2023年1-{month}月完成投资'
        self.ws['N5'].value = f'2024年1-{month}月预计完成投资'

        # 全部单元格字体颜色改为黑色
        font_black = Font(color='000000')

        for r in self.ws.iter_rows(min_row=6, max_col=21):
            for c in r:
                c.font = font_black


        self.wb.save(f'2024年{month}月-2024年能源以外工业项目投资情况统计表（5亿元以上项目）.xlsx')

        print(f'生成{month}月份的表格')



    def get_primary_col_values(self):
        values = []
        for row in self.ws.iter_rows(min_row=self.start_row, max_col=self.primary_col):
            values.append(row[self.primary_col - 1].value)
        no_none = [x for x in values if x is not None]
        return set(no_none)

    def get_delete_rows(self, primary_col_value):
        data_rows = []
        for i, row in enumerate(
            self.ws.iter_rows(min_row=self.start_row, max_col=self.max_col)
        ):
            if row[self.primary_col - 1].value != primary_col_value:
                data_rows.append(i + self.start_row)
        return data_rows

    def delete_rows(self, ws, to_delete_rows):
        to_delete_rows.reverse()
        for r in to_delete_rows:
            ws.delete_rows(r, 1)

    def split_by_primary_col(self, target_dir):
        primary_col_values = self.get_primary_col_values()
        print(primary_col_values)

        for key in primary_col_values:
            print(f'处理{key}表格')
            self.save_as(target_dir, key)

    def save_as(self, target_dir, key):
        wb = load_workbook(self.file_path)

        new_file = f"{target_dir}/{key}.xlsx"
        wb.save(new_file)

        new_wb = load_workbook(new_file)
        new_ws = new_wb.active

        merged_cells = new_ws.merged_cells
        mgcells = [m.coord for m in merged_cells if m.min_row >= self.start_row]
        for m in mgcells:
            new_ws.unmerge_cells(m)

        to_delete_rows = self.get_delete_rows(key)
        print(f"delete {to_delete_rows} rows for {key}")
        self.delete_rows(new_ws, to_delete_rows)

        new_ws.cell(3, 1).value = f"填报单位：{key}"
        new_wb.save(new_file)


def main():
    target_dir = "dst"
    if os.path.exists(target_dir):
        shutil.rmtree(target_dir)
    os.mkdir(target_dir)

    file_path = sys.argv[1]
    primary_col = 21
    start_row = 6
    max_col = 21

    month = sys.argv[2]

    ct = CollectTable(file_path, primary_col, start_row, max_col)

    ct.get_new_month(month)

    ct.split_by_primary_col(target_dir)


if __name__ == "__main__":
    main()
